import os
import pymongo
from loguru import logger
from pymongo.errors import DuplicateKeyError

from Api.retry_model import MyRetry
# from Config.config import Mongo_HOST, Mongo_USERNAME, Mongo_PASSWORD, Mongo_PORT, Mongo_DB


class MongoClient:
    def __init__(self, mongo_connect):
        try:
            if mongo_connect.get('user'):
                self.client = pymongo.MongoClient(
                    f'mongodb://{mongo_connect["user"]}:{mongo_connect["pwd"]}@{mongo_connect["host"]}:{mongo_connect["port"]}/{mongo_connect["db"]}'
                )
            else:
                self.client = pymongo.MongoClient(f'mongodb://{mongo_connect["host"]}:{mongo_connect["port"]}', )
            self.db = self.client[mongo_connect["db"]]

            logger.success("MongoDB连接成功！")
        except Exception as e:
            logger.error(f"MongoDB连接失败！{e}")
            os._exit(0)

    def aggregate(self, col, find_list, allowDiskUse=False):
        try:
            return True, self.db[col].aggregate(find_list, allowDiskUse=allowDiskUse)
        except Exception as e:
            logger.warning(e)
            return False, e

    def insert_document(self, col, document, Type=1):
        '''
        进行文档插入，插入成功返回 _id 否则返回False
        :param document: 当单一插入的时候是字典，否则是列表[{}]
        :return:
        '''
        if Type == 1:
            try:
                return True, self.db[col].insert_one(document)
            except DuplicateKeyError as e:
                logger.warning(f'{e}')
                return True, ''
            except Exception as e:
                logger.warning(e)
                return False, e
        elif Type == 3:
            for one_data in document:
                try:
                    self.db[col].insert_one(one_data)
                except DuplicateKeyError as e:
                    # logger.warning(f'{e}')
                    continue
                except Exception as e:
                    # logger.exception(e)
                    return (False, e)
            return True, ''
        else:
            try:
                return True, self.db[col].insert_many(document)
            except Exception as e:
                logger.error(f"写入数据失败，{e}")
                return False, e

    def update_document(self, col, query, update, Type=1, upsert=False):
        '''
        进行文档更新
        :param col: 集合
        :param query: 一个与要更新的文档匹配的查询
        :param update: 要应用的修改
        :return:
        '''
        if Type == 1:
            try:
                return True, self.db[col].update_one(query, update, upsert=upsert)
            except Exception as e:
                logger.error(f"更新失败，失败原因：{e}")
                return False, Exception
        else:
            try:
                return True, self.db[col].update_many(query, update, upsert=upsert)
            except Exception as e:
                logger.error(f"更新失败，失败原因：{e}")
                return False, e

    def update_different_filter_document(self, col, update_list):
        try:
            return True, self.db[col].bulk_write([pymongo.UpdateOne(i[0], i[1]) for i in update_list])
        except Exception as e:
            logger.error(f"更新失败，失败原因：{e}")
            return (False, Exception)

    def find_document(self, col, query, *args, **kwargs):
        '''
        进行文档查询
        :param query:   查询条件，必须是
        :param Tpye: 0/1    1查询单一文档，0查询所有文档
        :return:
        '''
        Type = kwargs.get('Type', 1)
        if kwargs.get('Type'):
            del kwargs['Type']
        try:
            if Type == 1:
                find_res = self.db[col].find_one(query, *args, **kwargs)
            else:
                find_res = self.db[col].find(query, *args, **kwargs)
            if find_res:
                return True, find_res
            else:
                return False, []
        except Exception as e:
            logger.warning(e)
            return False, []

    def remove_document(self, col, query, Type=1):
        if Type == 1:
            try:
                return True, self.db[col].delete_one(query)
            except Exception:
                logger.warning(Exception)
                return False, []
        else:
            try:
                return True, self.db[col].delete_many(query)
            except Exception:
                logger.warning(Exception)
                return False, []

    def success_update_db(self, col, md5):
        result = self.update_document(col, {"md5": md5}, {"$set": {"is_down": 1}})
        if result[0]:
            return True
        return False

    @MyRetry(times=3,  return_msg="更新mongodb 数据失败！")
    def faild_update_db(self, col, md5, is_down=1):
        result = self.update_document(col, {"md5": md5}, {"$set": {"is_down": is_down}})
        if result[0]:
            return True
        return False

    def count_documents_in_time_range(self, table_name, old_date, new_date):
        col = self.db[table_name]
        query = {
            'create_time': {
                '$gte': old_date,
                '$lte': new_date
            }
        }
        total_count = col.count_documents(query)
        return total_count


Mongo_HOST = "192.168.44.131"
Mongo_USERNAME = "ChineseAndEnglishBook_admin"
Mongo_DB = "ChineseAndEnglishBook"
Mongo_PASSWORD = "mZJy3XnbV7vwtKsY"
Mongo_PORT = 27017
mongo_connect = {
    "host": Mongo_HOST,
    "user": Mongo_USERNAME,
    "pwd": Mongo_PASSWORD,
    'port': Mongo_PORT,
    "db": Mongo_DB,
}
mongo_client = MongoClient(mongo_connect)

MONGO_HOST2 = "192.168.44.131"
MONGO_USER2 = "Book_admin"
MONGO_DB2 = "Book"
MONGO_PW2 = "QtO8I6weJV2KXrrA"
MONGO_PORT2 = 27017
mongo_connect2 = {
    "host": MONGO_HOST2,
    "user": MONGO_USER2,
    "pwd": MONGO_PW2,
    'port': MONGO_PORT2,
    "db": MONGO_DB2,
}
mongo_client2 = MongoClient(mongo_connect2)

